{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Variational Equations With the Chain Rule\n",
    "\n",
    "For a complete introduction to variational equations, please read the paper by Rein and Tamayo (2016).\n",
    "\n",
    "Variational equations can be used to calculate derivatives in an $N$-body simulation. More specifically, given a set of initial conditions $\\alpha_i$ and a set of variables at the end of the simulation $v_k$, we can calculate all first order derivatives\n",
    "$$\\frac{\\partial v_k}{\\partial \\alpha_i}$$\n",
    "as well as all second order derivates\n",
    "$$\\frac{\\partial^2 v_k}{\\partial \\alpha_i\\partial \\alpha_j}$$\n",
    "\n",
    "For this tutorial, we work with a two planet system. \n",
    "\n",
    "We first chose the semi-major axis $a$ of the outer planet as an initial condition (this is our $\\alpha_i$). At the end of the simulation we output the velocity of the star in the $x$ direction (this is our $v_k$). \n",
    "\n",
    "To do that, let us first import REBOUND and numpy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import rebound\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following function takes $a$ as a parameter, then integrates the two planet system and returns the velocity of the star at the end of the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_vx(a):\n",
    "    sim = rebound.Simulation()\n",
    "    sim.add(m=1.) # star\n",
    "    sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet\n",
    "    sim.add(primary=sim.particles[0],m=1e-3, a=a) # outer planet\n",
    "    sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits\n",
    "    return sim.particles[0].vx  # return star's velocity in the x direction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0004924175842478658"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "calculate_vx(a=1.5) # initial semi-major axis of the outer planet is 1.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we run the simulation again, with a different initial $a$, we get a different velocity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.000750246684761206"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "calculate_vx(a=1.51) # initial semi-major axis of the outer planet is 1.51"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could now run many different simulations to map out the parameter space. This is a very simple example of a typical use case: the fitting of a radial velocity datapoint. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, we can be smarter than simple running an almost identical simulation over and over again by using variational equations. These will allow us to calculate the *derivate* of the stellar velocity at the end of the simulation. We can take derivative with respect to any of the initial conditions, i.e. a particle's mass, semi-major axis, x-coordinate, etc. Here, we want to take the derivative with respect to the semi-major axis of the outer planet. The following function does exactly that:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_vx_derivative(a):\n",
    "    sim = rebound.Simulation()\n",
    "    sim.add(m=1.) # star\n",
    "    sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet\n",
    "    sim.add(primary=sim.particles[0],m=1e-3, a=a) # outer planet\n",
    "    \n",
    "    v1 = sim.add_variation()    # add a set of variational particles\n",
    "    v1.vary(2,\"a\")              # initialize the variational particles \n",
    "    \n",
    "    sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits\n",
    "    return sim.particles[0].vx, v1.particles[0].vx # return star's velocity and its derivative"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the two new functions. `sim.add_variation()` adds a set of variational particles to the simulation. All variational particles are by default initialized to zero. We use the `vary()` function to initialize them to a variation that we are interested in. Here, we initialize the variational particles corresponding to a change in the semi-major axis, $a$, of the particle with index 2 (the outer planet). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.0004924175842478302, 0.026958628196580445)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "calculate_vx_derivative(a=1.5) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the derivative to construct a Taylor series expansion of the velocity around $a_0=1.5$:\n",
    "$$v(a) \\approx v(a_0) + (a-a_0) \\frac{\\partial v}{\\partial a}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.000762003866214\n"
     ]
    }
   ],
   "source": [
    "a0=1.5\n",
    "va0, dva0 = calculate_vx_derivative(a=a0) \n",
    "def v(a):\n",
    "    return va0 + (a-a0)*dva0\n",
    "print(v(1.51))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare this value with the explicitly calculate one above. They are almost the same! But we can do even better, by using second order variational equations to calculate second order derivatives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def calculate_vx_derivative_2ndorder(a):\n",
    "    sim = rebound.Simulation()\n",
    "    sim.add(m=1.) # star\n",
    "    sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet\n",
    "    sim.add(primary=sim.particles[0],m=1e-3, a=a) # outer planet\n",
    "    \n",
    "    v1 = sim.add_variation()   \n",
    "    v1.vary(2,\"a\")             \n",
    "    \n",
    "    # The following lines add and initialize second order variational particles\n",
    "    v2 = sim.add_variation(order=2, first_order=v1)   \n",
    "    v2.vary(2,\"a\")  \n",
    "    \n",
    "    sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits\n",
    "    # return star's velocity and its first and second derivatives\n",
    "    return sim.particles[0].vx, v1.particles[0].vx, v2.particles[0].vx "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using a Taylor series expansion to second order gives a better estimate of `v(1.51)`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.000755071182773\n"
     ]
    }
   ],
   "source": [
    "a0=1.5\n",
    "va0, dva0, ddva0 = calculate_vx_derivative_2ndorder(a=a0) \n",
    "def v(a):\n",
    "    return va0 + (a-a0)*dva0 + 0.5*(a-a0)**2*ddva0\n",
    "print(v(1.51))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Now that we know how to calculate first and second order derivates of positions and velocities of particles, we can simply use the chain rule to calculate more complicated derivates. For example, instead of the velocity $v_x$, you might be interested in the quantity $w\\equiv(v_x - c)^2$ where $c$ is a constant. This is something that typically appears in a $\\chi^2$ fit. The chain rule gives us:\n",
    "$$ \\frac{\\partial w}{\\partial a} = 2 \\cdot (v_x-c)\\cdot \\frac{\\partial v_x}{\\partial a}$$\n",
    "The variational equations provide the $\\frac{\\partial v_x}{\\partial a}$ part, the *ordinary* particles provide $v_x$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1.039395710603212, -0.05496905171588172)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def calculate_w_derivative(a):\n",
    "    sim = rebound.Simulation()\n",
    "    sim.add(m=1.) # star\n",
    "    sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet\n",
    "    sim.add(primary=sim.particles[0],m=1e-3, a=a) # outer planet\n",
    "    \n",
    "    v1 = sim.add_variation()    # add a set of variational particles\n",
    "    v1.vary(2,\"a\")              # initialize the variational particles \n",
    "    \n",
    "    sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits\n",
    "    \n",
    "    c = 1.02 # some constant\n",
    "    w = (sim.particles[0].vx-c)**2\n",
    "    dwda = 2.*v1.particles[0].vx * (sim.particles[0].vx-c)\n",
    "    \n",
    "    return w, dwda # return w and its derivative\n",
    "calculate_w_derivative(1.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, you can also use the chain rule to vary initial conditions of particles in a way that is not supported by REBOUND by default. For example, suppose you want to work in some fancy coordinate system, using $h\\equiv e\\sin(\\omega)$ and $k\\equiv e \\cos(\\omega)$ variables instead of $e$ and $\\omega$. You might want to do that because $h$ and $k$ variables are often better behaved near $e\\sim0$. In that case the chain rule gives us:\n",
    "$$\\frac{\\partial p(e(h, k), \\omega(h, k))}{\\partial h} = \\frac{\\partial p}{\\partial e}\\frac{\\partial  e}{\\partial h} + \\frac{\\partial p}{\\partial \\omega}\\frac{\\partial \\omega}{\\partial h}$$\n",
    "where $p$ is any of the particles initial coordinates. In our case the derivates of $e$ and $\\omega$ with respect to $h$ are:\n",
    "$$\\frac{\\partial \\omega}{\\partial h} = -\\frac{k}{e^2}\\quad\\text{and}\\quad \\frac{\\partial e}{\\partial h} = \\frac{h}{e}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With REBOUND, you can easily implement this. The following function calculates the derivate of the star's velocity with respect to the outer planet's $h$ variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-0.0006022810748296454, 0.002107215810994136)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def calculate_vx_derivative_h():\n",
    "    h, k = 0.1, 0.2\n",
    "    e = float(np.sqrt(h**2+k**2))\n",
    "    omega = np.arctan2(k,h)\n",
    "    sim = rebound.Simulation()\n",
    "    sim.add(m=1.) # star\n",
    "    sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet\n",
    "    sim.add(primary=sim.particles[0],m=1e-3, a=1.5, e=e, omega=omega) # outer planet\n",
    "    \n",
    "    v1 = sim.add_variation()   \n",
    "    dpde     = rebound.Particle(simulation=sim, particle=sim.particles[2], variation=\"e\")\n",
    "    dpdomega = rebound.Particle(simulation=sim, particle=sim.particles[2], m=1e-3, a=1.5, e=e, omega=omega, variation=\"omega\")\n",
    "    v1.particles[2] = h/e * dpde - k/(e*e) * dpdomega\n",
    "        \n",
    "    sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits\n",
    "    # return star's velocity and its first derivatives\n",
    "    return sim.particles[0].vx, v1.particles[0].vx\n",
    "calculate_vx_derivative_h()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that in the above function, there are expressions such as `h/e * dpde`. `h/e` is just a number, but `dpde` is actually a particle structure. REBOUND multiplies each cartesian component of that particle with the number `h/e`. Similarly, the particles are subtracted componentwise when using the `-` operator.\n",
    "\n",
    "We can use the `v1.particles[i] = ...` syntax to directly set a variational particle's initial conditions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
